from pxr import Usd, UsdGeom, Gf
from queue import Queue


class USD_Analyzer():

    def __init__(self):
        pass


    """
    构建字典dic
    """
    def buildDic(self, usd_file):
        # 打开usd文件
        stage = Usd.Stage.Open(usd_file)
        root_prim = stage.GetPseudoRoot()
        blankDic = {}   # 字典

        def traversePrim(prim, dic):
            # 进行BFS遍历
            # 获取存入数据库的字典
            # 每一个prim节点的key为name， content为[childList, path, bboxList]
            queue = Queue()  # prim队列
            queue.put(prim)

            while (queue.qsize() > 0):
                curPrim = queue.get()
                primName = curPrim.GetName()
                # 将prim的属性list存入字典
                primPath = str(curPrim.GetPath())   # prim路径
                # 获取子name列表
                childList = []
                # 获取子bboxRange列表
                bboxList = []
                # 计算bboxRange
                bboxRange = self.compute_bbox(curPrim)


                # 注意这里是实例化后的对象
                for childPrim in curPrim.GetChildren():
                    # 不加进队列，不再细分
                    if (childPrim.IsInstance()):
                        childName = childPrim.GetName()
                        childList.append(childName)
                        childPath = str(childPrim.GetPath())   # prim路径
                        # 计算bboxRange
                        childBboxRange = self.compute_bbox(childPrim)
                        # 将Range3d顶点转换为tuple类型，字典格式存储
                        minPoint = tuple(childBboxRange.GetMin())
                        maxPoint = tuple(childBboxRange.GetMax())
                        bboxDic = {
                            "minPoint": minPoint,
                            "maxPoint": maxPoint
                        }
                        bboxList.append(bboxDic)

                        # 属性列表attributeDic
                        attributeDic = {}
                        # variantListDic为变体类型列表的字典，注意与variantDic区分
                        variantListDic = []

                        ###---------by tuo 6.27
                        # 打开usd_file并读取为字符串
                        file = open(usd_file, 'r')


                        file_content = file.read()
                        # 获取childname在文件中的位置
                        index = file_content.find(childName)
                        asset_name = None
                        if(index != -1):
                            index += len(childName)
                            # 从这后面开始找到 add references
                            sub = file_content[index:]
                            start = sub.find('references')
                            # 从start以后获取第三个字符串数组（通过空格分割
                            path = sub[start:].split() #去掉前后两个@字符
                            s = path[2]
                            p = s.split('/')
                            asset_name = p[2]


                            ###---------by tuo 6.27
                        file.close()


                        # 获取prim属性，以字典形式存放

                        # 注意！childAttribute_xformOp_translate
                        # childAttribute_xformOp_rotateXYZ
                        # childAttribute_xformOp_scale
                        # 要转换为tuple！！

                        # xformOpOrder，变换操作顺序，类型：list
                        childAttribute_xformOpOrder = childPrim.GetAttribute('xformOpOrder')
                        childAttribute_xformOpOrder = childAttribute_xformOpOrder.Get()
                        if childAttribute_xformOpOrder != None:
                            childAttribute_xformOpOrder = list(childAttribute_xformOpOrder)
                            # print(type(childAttribute_xformOpOrder))
                            attributeDic['xformOpOrder'] = childAttribute_xformOpOrder


                        # xformOp:translate，平移变换，类型：tuple
                        childAttribute_xformOp_translate = childPrim.GetAttribute('xformOp:translate')
                        childAttribute_xformOp_translate = childAttribute_xformOp_translate.Get()
                        if childAttribute_xformOp_translate != None:
                            childAttribute_xformOp_translate = tuple(childAttribute_xformOp_translate)
                            # print(type(childAttribute_xformOp_translate))
                            attributeDic['xformOp:translate'] = childAttribute_xformOp_translate


                        # xformOp:rotateX，绕x轴旋转，类型：float
                        childAttribute_xformOp_rotateX = childPrim.GetAttribute('xformOp:rotateX')
                        childAttribute_xformOp_rotateX = childAttribute_xformOp_rotateX.Get()
                        if childAttribute_xformOp_rotateX != None:
                            # print(type(childAttribute_xformOp_rotateX))
                            attributeDic['xformOp:rotateX'] = childAttribute_xformOp_rotateX


                        # xformOp:rotateY，绕y轴旋转，类型：float
                        childAttribute_xformOp_rotateY = childPrim.GetAttribute('xformOp:rotateY')
                        childAttribute_xformOp_rotateY = childAttribute_xformOp_rotateY.Get()
                        if childAttribute_xformOp_rotateY != None:
                            # print(type(childAttribute_xformOp_rotateY))
                            attributeDic['xformOp:rotateY'] = childAttribute_xformOp_rotateY


                        # xformOp:rotateZ，绕z轴旋转，类型：float
                        childAttribute_xformOp_rotateZ = childPrim.GetAttribute('xformOp:rotateZ')
                        childAttribute_xformOp_rotateZ = childAttribute_xformOp_rotateZ.Get()
                        if childAttribute_xformOp_rotateZ != None:
                            # print(type(childAttribute_xformOp_rotateZ))
                            attributeDic['xformOp:rotateZ'] = childAttribute_xformOp_rotateZ


                        # xformOp:rotateXYZ，绕xyz轴旋转，类型：tuple
                        childAttribute_xformOp_rotateXYZ = childPrim.GetAttribute('xformOp:rotateXYZ')
                        childAttribute_xformOp_rotateXYZ = childAttribute_xformOp_rotateXYZ.Get()
                        if childAttribute_xformOp_rotateXYZ != None:
                            childAttribute_xformOp_rotateXYZ = tuple(childAttribute_xformOp_rotateXYZ)
                            # print(type(childAttribute_xformOp_rotateXYZ))
                            attributeDic['xformOp:rotateXYZ'] = childAttribute_xformOp_rotateXYZ


                        # xformOp:scale，尺寸缩放，类型：tuple
                        childAttribute_xformOp_scale = childPrim.GetAttribute('xformOp:scale')
                        childAttribute_xformOp_scale = childAttribute_xformOp_scale.Get()
                        if childAttribute_xformOp_scale != None:
                            childAttribute_xformOp_scale = tuple(childAttribute_xformOp_scale)
                            # print(type(childAttribute_xformOp_scale))
                            attributeDic['xformOp:scale'] = childAttribute_xformOp_scale


                        # print(childAttribute_xformOpOrder)
                        # print(childAttribute_xformOp_translate)
                        # print(childAttribute_xformOp_rotateX)
                        # print(childAttribute_xformOp_rotateY)
                        # print(childAttribute_xformOp_rotateZ)
                        # print(childAttribute_xformOp_rotateXYZ)
                        # print(childAttribute_xformOp_scale)


                        # 获取当前prim的当前变体选择，以及所有变体可选内容
                        # 如果该prim拥有变体
                        if childPrim.HasVariantSets() == True:
                            variantSets = childPrim.GetVariantSets()
                            # print(" ")
                            # print("prim名字")
                            # print(childName)

                            # variantDic为variant的字典
                            for variantSetName in variantSets.GetNames():
                                # print(variantSetName)
                                variantDic = {}

                                # variantSet的名称
                                variantDic["variantSetName"] = variantSetName

                                # variantSet变体的当前选择
                                curVariantSelection = variantSets.GetVariantSelection(variantSetName)
                                variantDic["curVariantSelection"] = curVariantSelection
                                # print(curVariantSelection)

                                # variantSet变体的所有选择
                                variants = variantSets.GetVariantSet(variantSetName).GetVariantNames()
                                # 所有变体内容以列表形式存储
                                variantList = []
                                # print("所有变体可选内容：")
                                for variantName in variants:
                                    #  print(variantName)
                                    variantList.append(variantName)
                                variantDic["variantList"] = variantList

                                # 该类型prim是以reference指向的资产来衡量的
                                # 如果对于该prim的curVariantSelection为默认Default，则该条prim不具有这条变体
                                # 但是对于其reference指向的资产，拥有该条目，因此应该插入数据库
                                variantListDic.append(variantDic)

                                # if curVariantSelection != "Default":
                                #     variantListDic.append(variantDic)
                                # else:
                                #     continue

                            # print(variantListDic)

                        # prim没有变体
                        else:
                            variantListDic = None


                        # 叶子节点：
                        # 1. 无childList；（没有孩子prim）
                        # 2. 有自己作为孩子prim的primPath（这里起名叫childPath）；（在场景树中的路径）
                        # 3. 有bboxDic；（有自己作为孩子prim的bbox，里面存放的是minPoint和maxPoint）
                        # 4. 有assetName；（其下存放reference指向相对应的assets资产）
                        # 5. 有attributeDic；（其下存放关于属性变换操作的顺序，以及存放各个属性的值）
                        # 6. 有variantDic。 （其下存放关于当前变体选择，以及所有的变体可选内容）
                        dic[childName] = [None, childPath, bboxDic, asset_name, attributeDic, variantListDic]
                        continue

                    # 如果该孩子prim可再分
                    else:
                        # 获取子prim的名字
                        childName = childPrim.GetName()
                        childList.append(childName)
                        queue.put(childPrim)
                        # 计算bboxRange
                        childBboxRange = self.compute_bbox(childPrim)
                        # 将Range3d顶点转换为tuple类型，字典格式存储
                        minPoint = tuple(childBboxRange.GetMin())
                        maxPoint = tuple(childBboxRange.GetMax())
                        bboxDic = {
                            "minPoint": minPoint,
                            "maxPoint": maxPoint
                        }
                        bboxList.append(bboxDic)
                        continue

                # 中间节点：
                # 1. 有childList；（孩子prim的list）
                # 2. 有primPath；（在场景树中的路径）
                # 3. 有bboxList；（孩子prim的bbox的list）
                # 4. 没有assetName；（其下没有reference指向assets资产）
                # 5. 没有attributeDic；（中间节点没有属性）
                # 6. 没有variantDic。（中间节点没有变体）
                dic[primName] = [childList, primPath, bboxList, None, None, None]

            return dic


        dic = traversePrim(root_prim, blankDic)

        return dic


    """
    计算一个prim的bbox包围盒
    """
    def compute_bbox(self, prim):
        imageable = UsdGeom.Imageable(prim)
        time = Usd.TimeCode.Default()  # 计算bounding box的时间
        bound = imageable.ComputeWorldBound(time, UsdGeom.Tokens.default_)
        bound_range = bound.ComputeAlignedBox()
        return bound_range



